{
 "cells": [
  {
   "cell_type": "markdown",
   "id": "overview",
   "metadata": {},
   "source": [
    "# Strands Agents with AgentCore Memory (Short-Term Memory)\n",
    "\n",
    "\n",
    "## Introduction\n",
    "\n",
    "This tutorial demonstrates how to build a **personal agent** using Strands agents with AgentCore **short-term memory** (Raw events). The agent remembers recent conversations in the session using `get_last_k_turns` and can continue conversations seamlessly when user returns.\n",
    "\n",
    "\n",
    "### Tutorial Details\n",
    "\n",
    "| Information         | Details                                                                          |\n",
    "|:--------------------|:---------------------------------------------------------------------------------|\n",
    "| Tutorial type       | Short Term Conversational                                                        |\n",
    "| Agent type          | Personal Agent                                                                   |\n",
    "| Agentic Framework   | Strands Agents                                                                   |\n",
    "| LLM model           | Anthropic Claude Sonnet 3.7                                                      |\n",
    "| Tutorial components | AgentCore Short-term Memory, AgentInitializedEvent and MessageAddedEvent hooks   |\n",
    "| Example complexity  | Beginner                                                                         |\n",
    "\n",
    "You'll learn to:\n",
    "- Use short-term memory for conversation continuity\n",
    "- Retrieve last K conversation turns\n",
    "- Web search tool for real-time information\n",
    "- Initialize agents with conversation history\n",
    "\n",
    "## Architecture\n",
    "<div style=\"text-align:left\">\n",
    "    <img src=\"architecture.png\" width=\"65%\" />\n",
    "</div>\n",
    "\n",
    "## Prerequisites\n",
    "\n",
    "- Python 3.10+\n",
    "- AWS credentials with AgentCore Memory permissions\n",
    "- AgentCore Memory role ARN\n",
    "- Access to Amazon Bedrock models\n",
    "\n",
    "Let's get started by setting up our environment!"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "setup",
   "metadata": {},
   "source": [
    "## Step 1: Setup and Imports"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "641512ed",
   "metadata": {},
   "outputs": [],
   "source": [
    "!pip install -qr requirements.txt"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "setup_code",
   "metadata": {},
   "outputs": [],
   "source": [
    "import logging\n",
    "from datetime import datetime\n",
    "\n",
    "# Setup\n",
    "logging.basicConfig(level=logging.INFO)\n",
    "logger = logging.getLogger(\"personal-agent\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "imports",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Imports\n",
    "import os\n",
    "from strands import Agent, tool\n",
    "from strands.hooks import AgentInitializedEvent, HookProvider, HookRegistry, MessageAddedEvent\n",
    "from bedrock_agentcore.memory import MemoryClient\n",
    "\n",
    "# Configuration\n",
    "REGION = os.getenv('AWS_REGION', 'us-west-2') # AWS region for the agent\n",
    "ACTOR_ID = \"user_123\" # It can be any unique identifier (AgentID, User ID, etc.)\n",
    "SESSION_ID = \"personal_session_001\" # Unique session identifier"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "step2",
   "metadata": {},
   "source": [
    "## Step 2: Web Search Tool\n",
    "\n",
    "First, let's create a simple web search tool for the agent."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "web_search",
   "metadata": {},
   "outputs": [],
   "source": [
    "from ddgs.exceptions import DDGSException, RatelimitException\n",
    "from ddgs import DDGS\n",
    "\n",
    "@tool\n",
    "def websearch(keywords: str, region: str = \"us-en\", max_results: int = 5) -> str:\n",
    "    \"\"\"Search the web for updated information.\n",
    "    \n",
    "    Args:\n",
    "        keywords (str): The search query keywords.\n",
    "        region (str): The search region: wt-wt, us-en, uk-en, ru-ru, etc..\n",
    "        max_results (int | None): The maximum number of results to return.\n",
    "    Returns:\n",
    "        List of dictionaries with search results.\n",
    "    \n",
    "    \"\"\"\n",
    "    try:\n",
    "        results = DDGS().text(keywords, region=region, max_results=max_results)\n",
    "        return results if results else \"No results found.\"\n",
    "    except RatelimitException:\n",
    "        return \"Rate limit reached. Please try again later.\"\n",
    "    except DDGSException as e:\n",
    "        return f\"Search error: {e}\"\n",
    "    except Exception as e:\n",
    "        return f\"Search error: {str(e)}\"\n",
    "\n",
    "logger.info(\"✅ Web search tool ready\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "step3",
   "metadata": {},
   "source": [
    "## Step 3: Create Memory Resource\n",
    "For short-term memory, we create a memory resource without any strategies. This stores raw conversation turns that can be retrieved with `get_last_k_turns`.\n"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "create_memory",
   "metadata": {},
   "outputs": [],
   "source": [
    "from botocore.exceptions import ClientError\n",
    "\n",
    "# Initialize Memory Client\n",
    "client = MemoryClient(region_name=REGION)\n",
    "memory_name = \"PersonalAgentMemory\"\n",
    "\n",
    "try:\n",
    "    # Create memory resource without strategies (thus only access to short-term memory)\n",
    "    memory = client.create_memory_and_wait(\n",
    "        name=memory_name,\n",
    "        strategies=[],  # No strategies for short-term memory\n",
    "        description=\"Short-term memory for personal agent\",\n",
    "        event_expiry_days=7, # Retention period for short-term memory. This can be upto 365 days.\n",
    "    )\n",
    "    memory_id = memory['id']\n",
    "    logger.info(f\"✅ Created memory: {memory_id}\")\n",
    "except ClientError as e:\n",
    "    logger.info(f\"❌ ERROR: {e}\")\n",
    "    if e.response['Error']['Code'] == 'ValidationException' and \"already exists\" in str(e):\n",
    "        # If memory already exists, retrieve its ID\n",
    "        memories = client.list_memories()\n",
    "        memory_id = next((m['id'] for m in memories if m['id'].startswith(memory_name)), None)\n",
    "        logger.info(f\"Memory already exists. Using existing memory ID: {memory_id}\")\n",
    "except Exception as e:\n",
    "    # Show any errors during memory creation\n",
    "    logger.error(f\"❌ ERROR: {e}\")\n",
    "    import traceback\n",
    "    traceback.print_exc()\n",
    "    # Cleanup on error - delete the memory if it was partially created\n",
    "    if memory_id:\n",
    "        try:\n",
    "            client.delete_memory_and_wait(memory_id=memory_id)\n",
    "            logger.info(f\"Cleaned up memory: {memory_id}\")\n",
    "        except Exception as cleanup_error:\n",
    "            logger.error(f\"Failed to clean up memory: {cleanup_error}\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "step4",
   "metadata": {},
   "source": [
    "## Step 4: Memory Hook\n",
    "\n",
    "This step defines our custom `MemoryHookProvider` class that automates memory operations. Hooks are special functions that run at specific points in an agent's execution lifecycle. The memory hook we're creating serves two primary functions:\n",
    "1. **To load recent conversation**: We use the `AgentInitializedEvent` hook will automatically load recent conversation history when the agent is initialized.\n",
    "2. **To store the last message**: Stores new conversational message.\n",
    "\n",
    "This creates a seamless memory experience without manual management."
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "memory_hook",
   "metadata": {},
   "outputs": [],
   "source": [
    "class MemoryHookProvider(HookProvider):\n",
    "    def __init__(self, memory_client: MemoryClient, memory_id: str, actor_id: str, session_id: str):\n",
    "        self.memory_client = memory_client\n",
    "        self.memory_id = memory_id\n",
    "        self.actor_id = actor_id\n",
    "        self.session_id = session_id\n",
    "    \n",
    "    def on_agent_initialized(self, event: AgentInitializedEvent):\n",
    "        \"\"\"Load recent conversation history when agent starts\"\"\"\n",
    "        try:\n",
    "            # Load the last 5 conversation turns from memory\n",
    "            recent_turns = self.memory_client.get_last_k_turns(\n",
    "                memory_id=self.memory_id,\n",
    "                actor_id=self.actor_id,\n",
    "                session_id=self.session_id,\n",
    "                k=5\n",
    "            )\n",
    "            \n",
    "            if recent_turns:\n",
    "                # Format conversation history for context\n",
    "                context_messages = []\n",
    "                for turn in recent_turns:\n",
    "                    for message in turn:\n",
    "                        role = message['role']\n",
    "                        content = message['content']['text']\n",
    "                        context_messages.append(f\"{role}: {content}\")\n",
    "                \n",
    "                context = \"\\n\".join(context_messages)\n",
    "                # Add context to agent's system prompt.\n",
    "                event.agent.system_prompt += f\"\\n\\nRecent conversation:\\n{context}\"\n",
    "                logger.info(f\"✅ Loaded {len(recent_turns)} conversation turns\")\n",
    "                \n",
    "        except Exception as e:\n",
    "            logger.error(f\"Memory load error: {e}\")\n",
    "    \n",
    "    def on_message_added(self, event: MessageAddedEvent):\n",
    "        \"\"\"Store messages in memory\"\"\"\n",
    "        messages = event.agent.messages\n",
    "        try:\n",
    "            self.memory_client.create_event(\n",
    "                memory_id=self.memory_id,\n",
    "                actor_id=self.actor_id,\n",
    "                session_id=self.session_id,\n",
    "                messages=[(messages[-1][\"content\"][0][\"text\"], messages[-1][\"role\"])]\n",
    "            )\n",
    "        except Exception as e:\n",
    "            logger.error(f\"Memory save error: {e}\")\n",
    "    \n",
    "    def register_hooks(self, registry: HookRegistry):\n",
    "        # Register memory hooks\n",
    "        registry.add_callback(MessageAddedEvent, self.on_message_added)\n",
    "        registry.add_callback(AgentInitializedEvent, self.on_agent_initialized)"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "step5",
   "metadata": {},
   "source": [
    "## Step 5: Create Personal Agent with Web Search"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "create_agent",
   "metadata": {},
   "outputs": [],
   "source": [
    "def create_personal_agent():\n",
    "    \"\"\"Create personal agent with memory and web search\"\"\"\n",
    "    agent = Agent(\n",
    "        name=\"PersonalAssistant\",\n",
    "        system_prompt=f\"\"\"You are a helpful personal assistant with web search capabilities.\n",
    "        \n",
    "        You can help with:\n",
    "        - General questions and information lookup\n",
    "        - Web searches for current information\n",
    "        - Personal task management\n",
    "        \n",
    "        When you need current information, use the websearch function.\n",
    "        Today's date: {datetime.today().strftime('%Y-%m-%d')}\n",
    "        Be friendly and professional.\"\"\",\n",
    "        hooks=[MemoryHookProvider(client, memory_id, ACTOR_ID, SESSION_ID)],\n",
    "        tools=[websearch],\n",
    "    )\n",
    "    return agent\n",
    "\n",
    "# Create agent\n",
    "agent = create_personal_agent()\n",
    "logger.info(\"✅ Personal agent created with memory and web search\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "step6",
   "metadata": {},
   "source": [
    "#### Congratulations ! Your agent is ready ! :) \n",
    "## Lets test the Agent"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "first_session",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Test conversation with memory\n",
    "print(\"=== First Conversation ===\")\n",
    "print(f\"User: My name is Alex and I'm interested in learning about AI.\")\n",
    "print(f\"Agent: \", end=\"\")\n",
    "agent(\"My name is Alex and I'm interested in learning about AI.\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "second_message",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"User: Can you search for the latest AI trends in 2025?\")\n",
    "print(f\"Agent: \", end=\"\")\n",
    "agent(\"Can you search for the latest AI trends in 2025?\")"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "third_message",
   "metadata": {},
   "outputs": [],
   "source": [
    "print(f\"User: I'm particularly interested in machine learning applications.\")\n",
    "print(f\"Agent: \", end=\"\")\n",
    "agent(\"I'm particularly interested in machine learning applications.\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "step7",
   "metadata": {},
   "source": [
    "## Test Memory Continuity\n",
    "\n",
    "To test if our memory system is working correctly, we'll create a new instance of the agent and see if it can access the previously stored information:"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "agent_restart",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Create new agent instance (simulates user returning)\n",
    "print(\"=== User Returns - New Session ===\")\n",
    "new_agent = create_personal_agent()\n",
    "\n",
    "# Test memory continuity\n",
    "print(f\"User: What was my name again?\")\n",
    "print(f\"Agent: \", end=\"\")\n",
    "new_agent(\"What was my name again?\")\n",
    "\n",
    "print(f\"User: Can you search for more information about machine learning?\")\n",
    "print(f\"Agent: \", end=\"\")\n",
    "new_agent(\"Can you search for more information about machine learning?\")"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "step8",
   "metadata": {},
   "source": [
    "## View Stored Memory"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "verify_memory",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Check what's stored in memory\n",
    "print(\"=== Memory Contents ===\")\n",
    "recent_turns = client.get_last_k_turns(\n",
    "    memory_id=memory_id,\n",
    "    actor_id=ACTOR_ID,\n",
    "    session_id=SESSION_ID,\n",
    "    k=3 # Adjust k to see more or fewer turns\n",
    ")\n",
    "\n",
    "for i, turn in enumerate(recent_turns, 1):\n",
    "    print(f\"Turn {i}:\")\n",
    "    for message in turn:\n",
    "        role = message['role']\n",
    "        content = message['content']['text'][:100] + \"...\" if len(message['content']['text']) > 100 else message['content']['text']\n",
    "        print(f\"  {role}: {content}\")\n",
    "    print()"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "summary",
   "metadata": {},
   "source": [
    "## Summary\n",
    "\n",
    "This tutorial showed how to build a personal agent. You've learned:\n",
    "\n",
    "- Creating memory resources without strategies\n",
    "- Using `get_last_k_turns` for conversation history\n",
    "- Adding web search capabilities to agents\n",
    "- Implementing memory hooks for context loading\n",
    "\n",
    "**Next Steps:**\n",
    "- Add more sophisticated tools\n",
    "- Implement long-term memory strategies\n",
    "- Enhance search capabilities with multiple sources"
   ]
  },
  {
   "cell_type": "markdown",
   "id": "cleanup",
   "metadata": {},
   "source": [
    "## Cleanup (Optional)"
   ]
  },
  {
   "cell_type": "code",
   "execution_count": null,
   "id": "cleanup_code",
   "metadata": {},
   "outputs": [],
   "source": [
    "# Uncomment to delete memory resource\n",
    "# client.delete_memory_and_wait(memory_id)\n",
    "# logger.info(f\"✅ Deleted memory: {memory_id}\")"
   ]
  }
 ],
 "metadata": {
  "kernelspec": {
   "display_name": "conda_pytorch_p310",
   "language": "python",
   "name": "conda_pytorch_p310"
  },
  "language_info": {
   "codemirror_mode": {
    "name": "ipython",
    "version": 3
   },
   "file_extension": ".py",
   "mimetype": "text/x-python",
   "name": "python",
   "nbconvert_exporter": "python",
   "pygments_lexer": "ipython3",
   "version": "3.10.14"
  }
 },
 "nbformat": 4,
 "nbformat_minor": 5
}
